# -*- coding: utf-8 -*-
"""
Manim脚本生成器 - 使用LLM生成Manim动画脚本
"""

import re
from typing import Any, Dict, List, Optional


class ScriptGenerator:
    """Manim脚本生成器"""

    def __init__(self, provider):
        """
        初始化生成器

        Args:
            provider: LLM Provider实例
        """
        self.provider = provider

    def generate(self, scenes: List[Dict[str, Any]], captions: str, resolution: str, fps: int) -> Optional[str]:
        """
        生成Manim脚本

        Args:
            scenes: 场景列表
            captions: 字幕文本
            resolution: 分辨率
            fps: 帧率

        Returns:
            Python代码字符串
        """
        prompt = self._build_prompt(scenes, captions, resolution, fps)
        messages = [
            {"role": "system", "content": "你只需输出 Python 代码，不要任何解释。"},
            {"role": "user", "content": prompt},
        ]

        content = self.provider.chat(messages, temperature=0.3)
        if not content:
            return None

        # 移除代码围栏
        content = re.sub(r"^```(?:python)?\n|\n```$", "", content.strip())

        # 验证基本结构
        if "class Generated" in content or "class Generated(Scene)" in content:
            return content

        return None

    def _build_prompt(self, scenes: List[Dict[str, Any]], captions: str, resolution: str, fps: int) -> str:
        """构建提示词"""
        lines = []
        lines.append("你是一名 Manim 动画工程师。根据给定的场景信息与字幕，生成一个 Manim 脚本。")
        lines.append("要求：")
        lines.append("1. 必须输出完整的 Python 代码，且包含 `from manim import *` 和 `class Generated(Scene)`，在 `construct` 中实现。")
        lines.append("2. 禁止任何网络/文件写入操作，不要读取外部资源。")
        lines.append("3. 使用 Text/MathTex 等基础元素，演示逐场景的呈现与过渡。")
        lines.append(f"4. 目标视频参数：分辨率 {resolution}，FPS {fps}。")
        lines.append("")
        lines.append("场景信息(JSON 简述)：")
        for s in scenes:
            title = (s.get("title") or "").strip()
            narration = (s.get("narration") or "").strip()
            lines.append(f"- 标题: {title} / 台词: {narration}")

        if captions:
            lines.append("")
            lines.append("字幕(SRT)：")
            lines.append(captions)

        return "\n".join(lines)
